#!/usr/bin/env python
"""

    download_liftover_data.py
    [--log_file PATH]
    [--verbose]

"""

################################################################################
#
#   prepare_liftover_data
#
#
#   Copyright (c) 12/11/2009 Leo Goodstadt
#
#   Permission is hereby granted, free of charge, to any person obtaining a copy
#   of this software and associated documentation files (the "Software"), to deal
#   in the Software without restriction, including without limitation the rights
#   to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#   copies of the Software, and to permit persons to whom the Software is
#   furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
#   THE SOFTWARE.
#################################################################################

import sys, os, re, tempfile
from collections import defaultdict
from run_cmd import run_cmd
from parse_delimited import parse_delimited
from lg_program_logging import setup_std_logging, MESSAGE

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Functions


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
def initial_lowercase (s):
    """
    useful for UCSC species names
    """
    return s[0:1].lower() + s[1:]

#_________________________________________________________________________________________
#
#   wget_helper
#
#_________________________________________________________________________________________
def wget_helper (local_dir, remote_file, logger):

    # temp log file
    wget_log = tempfile.NamedTemporaryFile(delete=False, suffix = "log")
    wget_log_file = wget_log.name
    wget_log.close()

    cmd = ["wget",
           "--directory-prefix=%s" % local_dir,
           "--timestamping",
           remote_file,
           "2> " + wget_log_file
           ]
    cmd = " ".join(cmd)
    os.system(cmd)

    # forward wget log to logger
    for l in open(wget_log_file):
        logger.info(l.rstrip())
    os.unlink(wget_log_file)

#_________________________________________________________________________________________
#
#   download_liftover_files
#
#_________________________________________________________________________________________
def blank_str_if_null (d):
    if d == None:
        return ""
    return str(d)


def download_liftover_files(seed_ucsc_genome, liftover_species_file,
                            ucsc_golden_path_ftp, liftover_data_dir,
                            liftover_parameters_file, logger):
    """
    Download liftover files from UCSC

    Save results in lifover_parmeters_file

    """

    #      1) Get list of ucsc genomes specified in "additional.genomes."
    #
    #      2) Download all liftover files for the reference genome of the seed transcripts
    #         (They are relatively small, so might as well download them all)
    #
    #      3) Gunzip all *.gz in liftover directory
    #
    #      4) Save list of additional genome species
    #
    #      Nomenclature:
    #          genome      = ucsc name                                e.g. GalGal
    #          genome_lc   = ucsc name without initial capitalization e.g. galGal
    #          species     = latin name                               e.g. gallus_gallus



    #
    #   get list of ucsc genomes whose liftover loci we want
    #
    (indirect_liftovers_to_via,
     direct_liftovers_from_to,
     ucsc_genome_to_species,
     min_liftover_lengths,
     via_species     )= get_list_of_ucsc_genomes (logger, liftover_species_file, seed_ucsc_genome)




    liftover_re= re.compile("([A-Za-z]+)([0-9]+)To([A-Za-z]+)([0-9]+).over.chain.gz" )
    missing_liftover_files = []
    lift_to_genomes_filenames = dict()


    #
    #   Download all liftover files for the reference genome of the seed transcripts
    #       and get list of actually downloaded files
    #
    logger.log(MESSAGE, "Download resources for liftover between source and destination species")
    seed_ucsc_genome_lc     = initial_lowercase(seed_ucsc_genome)
    wget_helper (liftover_data_dir,
                 ucsc_golden_path_ftp +
                  '/%s/liftOver/*.over.chain.gz' % seed_ucsc_genome_lc, logger)
    logger.log(MESSAGE, "Check downloaded direct liftover files")
    get_list_of_downloaded_liftover_files (liftover_re, direct_liftovers_from_to,
                                           lift_to_genomes_filenames,
                                           missing_liftover_files,
                                           liftover_data_dir, logger)

    #
    #   Down indirect liftover files if a direct liftover link does not exist
    #       and get list of actually downloaded files
    #
    logger.log(MESSAGE, "Download resources for liftover which involve an additional intermediate species")
    indirect_liftovers_via_to= download_indirect_liftover_files (lift_to_genomes_filenames,
                                                                 indirect_liftovers_to_via,
                                                                 missing_liftover_files,
                                                                 ucsc_golden_path_ftp,
                                                                 liftover_data_dir,
                                                                 logger)
    logger.log(MESSAGE, "Check downloaded indirect liftover files")
    get_list_of_downloaded_liftover_files ( liftover_re, indirect_liftovers_via_to,
                                            lift_to_genomes_filenames,
                                            missing_liftover_files,
                                            liftover_data_dir, logger)





    #
    #   Gunzip all *.gz in liftover directory
    #
    logger.log(MESSAGE, "Gunzip all gzipped liftover files")
    gunzip_all_gzfiles_in_directory (liftover_data_dir, logger)



    #
    #   return liftover parameters (i.e. paths to actual liftover files and parameters)
    #
    logger.log(MESSAGE, "Write liftover parameters")
    liftover_parameters_file  = open(liftover_parameters_file, "w")
    liftover_parameters_file.write("\t".join(["species", "ucsc genome", "ucsc genome version",
                                            "liftover file1", "liftover file2", "via_species",
                                            "min_liftover_len1", "min_liftover_len2"])+ "\n")

    for to_genome in lift_to_genomes_filenames:
        to_genome_version, lift_file1, lift_file2 = lift_to_genomes_filenames[to_genome]
        to_species = ucsc_genome_to_species[to_genome]
        (min_liftover_len1,
         min_liftover_len2) = min_liftover_lengths[to_species]
        lift_file_path1 = os.path.join(liftover_data_dir,lift_file1)
        if lift_file2:
            lift_file_path2 = os.path.join(liftover_data_dir,lift_file2)
        else:
            lift_file_path2 = None
        data = [to_species, to_genome, to_genome_version, lift_file_path1,
                lift_file_path2, via_species[to_species], min_liftover_len1, min_liftover_len2]
        data = "\t".join(map(blank_str_if_null, data)) + "\n"
        liftover_parameters_file.write(data)


    if missing_liftover_files:
        raise Exception("Missing liftover files. See log for more details.")



#_________________________________________________________________________________________
#
#   1) get_list_of_ucsc_genomes
#
#_________________________________________________________________________________________
def get_list_of_ucsc_genomes (logger, liftover_species_file, seed_ucsc_genome):
    """
    Read list from ``additional.genomes``
    """
    ucsc_genomes                = set()
    direct_liftovers_from_to    = defaultdict(set)
    indirect_liftovers_to_via   = dict()
    ucsc_genome_to_species      = dict()
    min_liftover_lengths        = dict()

    # whether direct liftover or not
    via_species                 = dict()

    for (lift_to_genome,
         lift_to_species,
         min_liftover_len1,
         lift_via_genome,
         lift_via_species,
         min_liftover_len2,
        ) in parse_delimited (open(liftover_species_file),
                                               field_names = ["UCSC_name",
                                                              "species",
                                                              "min_liftover_len1",
                                                              "liftover_via_UCSC",
                                                              "liftover_via_species",
                                                              "min_liftover_len2",
                                                              ],
                                               header=True):

        min_liftover_lengths[lift_to_species] = [min_liftover_len1, min_liftover_len2]

        # don't need liftover when it is the same species
        # skip
        if lift_to_genome == seed_ucsc_genome.rstrip("0123456789"):
            continue

        ucsc_genome_to_species[lift_to_genome] = lift_to_species
        ucsc_genomes.add(lift_to_genome)
        via_species[lift_to_species] = None

        if lift_via_genome:
            ucsc_genomes.add(lift_via_genome)
            indirect_liftovers_to_via[lift_to_genome] = lift_via_genome
            ucsc_genome_to_species[lift_via_genome] = lift_via_species
            via_species[lift_to_species] = lift_via_species
        else:
            # N.B.
            #   The "from_genome" (dictionary key)   contains the genome version
            #   The "to_genome"   (dictionary value) does not contain the genome version
            #                                        Which is the most uptodate version
            #                                        is only known after download
            direct_liftovers_from_to[seed_ucsc_genome].add(lift_to_genome)

    logger.info("ucsc genomes = " + ", ".join(ucsc_genomes))

    return (indirect_liftovers_to_via, direct_liftovers_from_to,
            ucsc_genome_to_species, min_liftover_lengths, via_species)



#_________________________________________________________________________________________
#
#   2/4) get_list_of_downloaded_liftover_files
#
#_________________________________________________________________________________________
def get_list_of_downloaded_liftover_files (liftover_re, liftovers_from_to,
                                           lift_to_genomes_filenames,
                                           missing_liftover_files,
                                           liftover_data_dir, logger):

    for f in os.listdir(liftover_data_dir):
        re_match = liftover_re.search(f)
        if re_match:
            from_genome = re_match.group(1)[0].upper() + re_match.group(1)[1:]
            from_genome_version = re_match.group(2)
            to_genome = re_match.group(3)
            to_genome_version = re_match.group(4)

            # make sure these liftover files are the right genome
            # N.B. for liftovers_from_to
            #   The "from_genome" (dictionary key)   contains the genome version
            #   The "to_genome"   (dictionary value) does not contain the genome version
            #                                        Which is the most uptodate version
            #                                        is only known after download
            if from_genome + from_genome_version not in liftovers_from_to:
                continue
            if to_genome not in liftovers_from_to[from_genome + from_genome_version]:
                continue

            to_genome_version = int(to_genome_version)

            #
            # use latest version of genome (i.e. mm9 preferred to mm8)
            #
            if (to_genome not in lift_to_genomes_filenames or
                to_genome_version > lift_to_genomes_filenames[to_genome][1]):
                # indirect liftover
                if from_genome in lift_to_genomes_filenames:
                    lift_to_genomes_filenames[to_genome] = (to_genome_version,
                                                            lift_to_genomes_filenames[from_genome][1],
                                                            f[0:-3])
                # direct liftover
                else:
                    lift_to_genomes_filenames[to_genome] = (to_genome_version, f[0:-3], None)


    #
    #   see if any liftover download failed
    #
    for from_genome, to_genomes in liftovers_from_to.iteritems():
        for to_genome in to_genomes:
            if to_genome not in lift_to_genomes_filenames:
                logger.error("Missing liftover file for UCSC genome %s from %s" %
                             (to_genome, from_genome))
                missing_liftover_files.append(to_genome)
#



#_________________________________________________________________________________________
#
#   3) download_indirect_liftover_files
#
#_________________________________________________________________________________________
def download_indirect_liftover_files (lift_to_genomes_filenames, indirect_liftovers_to_via,
                                        missing_liftover_files, ucsc_golden_path_ftp,
                                        liftover_data_dir, logger):
    """
    Genomes which need an extra liftover hop to get there e.g. mm9->GalGal3->TaeGut2
    """
    indirect_liftovers_via_to = defaultdict(set)
    for lift_to_genome, lift_via_genome in indirect_liftovers_to_via.iteritems():
        if lift_via_genome not in lift_to_genomes_filenames:
            logger.error("Missing liftover file for UCSC species (%s) " % lift_via_genome +
                         "used as an liftover intermediary to reach %s" %
                         (lift_to_genome))
            missing_liftover_files.append(lift_via_genome)
            continue
        lift_via_genome_version = lift_to_genomes_filenames[lift_via_genome][0]
        # N.B.
        #   The "from_genome" (dictionary key)   contains the genome version
        #   The "to_genome"   (dictionary value) does not contain the genome version
        #                                        Which is the most uptodate version
        #                                        is only known after download
        indirect_liftovers_via_to[lift_via_genome + str(lift_via_genome_version)].add(lift_to_genome)
        lift_via_genome_lc = initial_lowercase(lift_via_genome) + str(lift_via_genome_version)
        liftover_file_name = "%sTo%s*.over.chain.gz" % (lift_via_genome_lc, lift_to_genome)
        wget_helper (liftover_data_dir, ucsc_golden_path_ftp +
                                    '/%s/liftOver/%s' % (lift_via_genome_lc, liftover_file_name),
                                    logger)
    return indirect_liftovers_via_to



#_________________________________________________________________________________________
#
#   5) gunzip_all_gzfiles_in_directory
#
#_________________________________________________________________________________________
def gunzip_all_gzfiles_in_directory (directory, logger):
    """Leaves gz files in place
    """
    all_gzipped_files = set(os.listdir(directory))
    for f in all_gzipped_files:
        if f[-3:] == ".gz" and f[:-3] not in all_gzipped_files:
            logger.debug("Gunzipping %s" % f)
            full_path = os.path.join(directory, f)
            run_cmd("gunzip %s -c " % full_path, "unzipping %s" % f, logger = logger,
                        stdout_file_name = full_path[:-3])
            #os.system("gunzip %s -c > %s" % (full_path, full_path[:-3]))


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Testing


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
if __name__ == '__main__':
    import logging, os, sys
    module_name = os.path.splitext(os.path.basename(sys.argv[0]))[0]

    logger = logging.getLogger(module_name)
    setup_std_logging(logger, "liftover.log", 5)

    def get_logger (logger_name, args):
        return logger

    import unittest
    class Test_prepare_liftover_data(unittest.TestCase):

        #       self.assertEqual(self.seq, range(10))
        #       self.assert_(element in self.seq)
        #       self.assertRaises(ValueError, random.sample, self.seq, 20)



        def test_function(self):
            """
                test
            """
            download_liftover_files(
                                    "Mm9",                                                                      # SEED_TRANSCRIPT_UCSC_GENOME
                                    "/net/cpp-group/Leo/inprogress/brain_stuff/data/ucsc/additional.genomes",   # ADDITIONAL_GENOMES
                                    "ftp://hgdownload.cse.ucsc.edu/goldenPath",                                 # UCSC_GOLDEN_PATH_FTP
                                    "/net/cpp-group/Leo/inprogress/brain_stuff/data/liftover_resources",        # LIFTOVER_DIR_ROOT
                                    "/net/cpp-group/Leo/inprogress/brain_stuff/data/ucsc/Mm9.liftover_data",    # DATA_PARAMETERS_FOR_LIFTOVER
                                    logger,
                                                        )

#
#   debug code not run if called as a module
#
if __name__ == '__main__':
    if sys.argv.count("--debug"):
        sys.argv.remove("--debug")
    unittest.main()




